import uuid
from datetime import datetime
from collections import deque


class HistoryMemory:
    """
    管理对话历史记录的类，支持会话 ID 和上下文检索。
    """

    def __init__(self, max_sessions=10, max_turns_per_session=20):
        """
        初始化 HistoryMemory 实例。

        Args:
            max_sessions (int): 保留的最大会话数。默认为 10。
            max_turns_per_session (int): 每个会话保留的最大对话轮数。默认为 20。
        """
        self.max_sessions = max_sessions
        self.max_turns_per_session = max_turns_per_session

        # 存储所有会话 {session_id: deque of turns}
        # 使用 OrderedDict 或简单 dict 来模拟 LRU，这里简化处理
        self.sessions = {}
        # 当前活动会话 ID
        self.current_session_id = None

    def _get_or_create_session(self, session_id=None):
        """获取现有会话或创建新会话"""
        if session_id is None:
            #生成128位的id
            session_id = str(uuid.uuid4())

        if session_id not in self.sessions:
            # 创建新的会话历史 创建双端队列 双端队列支持高效地在两端添加/删除元素
            self.sessions[session_id] = deque(maxlen=self.max_turns_per_session)
            # 简单的会话数量管理
            if len(self.sessions) > self.max_sessions:
                # 移除最早的一个会话 (简单实现)
                # iter指向最早的会话
                # next 取到第一个会话的ID
                first_key = next(iter(self.sessions))
                del self.sessions[first_key]

        self.current_session_id = session_id
        return self.sessions[session_id]

    def start_new_session(self):
        """开始一个新会话"""
        new_session_id = str(uuid.uuid4())
        self._get_or_create_session(new_session_id)
        return new_session_id

    def add_turn(self, user_input: str, agent_response: str, session_id=None):
        """
        添加一轮对话到指定或当前会话的历史记录中。

        Args:
            user_input (str): 用户的输入。
            agent_response (str): Agent 的回复。
            session_id (str, optional): 会话 ID。如果为 None，则使用当前会话。
        """
        session_history = self._get_or_create_session(session_id)
        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        turn = {
            "id": str(uuid.uuid4()),  # 为每轮对话分配唯一 ID
            "timestamp": timestamp,
            "user": user_input,
            "agent": agent_response
        }
        session_history.append(turn)
        # print(f"Added turn to session {self.current_session_id}: {turn}") # 调试用

    def get_history(self, session_id=None) -> list:
        """
        获取指定或当前会话的对话历史记录列表。

        Args:
            session_id (str, optional): 会话 ID。如果为 None，则使用当前会话。

        Returns:
            list: 包含对话历史的字典列表。
        """
        session_id = session_id or self.current_session_id
        # 先确保session_id有效，再检查它是否在已保存的会话列表中
        if session_id and session_id in self.sessions:
            #存在则返回该会话的历史记录（转换为列表形式）
            return list(self.sessions[session_id])
        #不存在则返回空列表（安全兜底）
        return []

    def get_context(self, session_id=None, last_n=5) -> str:
        """
        获取指定或当前会话的上下文字符串，用于构建 Prompt。
        格式化为易于模型理解的对话历史。

        Args:
            session_id (str, optional): 会话 ID。
            last_n (int): 包含最近的 N 轮对话。

        Returns:
            str: 格式化的上下文字符串。
        """
        history = self.get_history(session_id)
        # 取最近的 last_n 轮
        # history[-last_n:] 从列表末尾往前数 last_n 个位置，取到列表最后一个元素
        recent_history = history[-last_n:] if len(history) >= last_n else history

        context_parts = []
        for turn in recent_history:
            context_parts.append(f"用户: {turn['user']}")
            context_parts.append(f"助手: {turn['agent']}")

        return "\n".join(context_parts)

    def get_sessions_list(self) -> list:
        """
        获取所有会话的简要信息列表，用于左侧导航。
        按时间倒序排列（最新的在前）。

        Returns:
            list: 包含会话信息的字典列表。
                  格式: [{'id': '...', 'title': '...', 'timestamp': '...'}, ...]
        """
        session_list = []
        for sid, history in self.sessions.items():
            if history:
                # 使用第一轮对话的用户输入或时间作为标题
                first_turn = history[0]
                title = first_turn['user'][:30] + ('...' if len(first_turn['user']) > 30 else '')
                timestamp = first_turn['timestamp']
            else:
                title = "新会话"
                timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

            session_list.append({
                "id": sid,
                "title": title,
                "timestamp": timestamp
            })
        # 按时间戳倒序排序 (假设时间戳是可比较的字符串)
        # 注意：如果时间戳格式不一致或跨越时区，可能需要更复杂的排序
        session_list.sort(key=lambda x: x['timestamp'], reverse=True)
        return session_list

    def clear_session(self, session_id=None):
        """
        清空指定或当前会话的历史记录。

        Args:
            session_id (str, optional): 会话 ID。如果为 None，则清空当前会话。
        """
        session_id = session_id or self.current_session_id
        if session_id and session_id in self.sessions:
            self.sessions[session_id].clear()
            # print(f"Session {session_id} cleared.") # 调试用

    def delete_session(self, session_id):
        """
        删除指定的会话。

        Args:
            session_id (str): 要删除的会话 ID。
        """
        if session_id in self.sessions:
            del self.sessions[session_id]
            if self.current_session_id == session_id:
                self.current_session_id = None
            # print(f"Session {session_id} deleted.") # 调试用

    def clear_all_sessions(self):
        """
        清空所有会话（删除所有会话历史，重置当前会话ID）
        """
        # 清空所有会话字典
        self.sessions.clear()
        # 重置当前活动会话ID（无活跃会话）
        self.current_session_id = None
        print("✅ 所有会话已清空")  # 调试用，可删除

    def set_current_session(self, session_id):
        """
        设置当前活动的会话。

        Args:
            session_id (str): 要设置为当前的会话 ID。
        """
        if session_id in self.sessions:
            self.current_session_id = session_id
        else:
            # 如果会话不存在，可以选择创建或忽略
            # 这里选择忽略
            pass


# --- 示例用法 ---
if __name__ == "__main__":
    memory = HistoryMemory(max_sessions=3, max_turns_per_session=5)

    # 开始新会话 1
    session1_id = memory.start_new_session()
    memory.add_turn("你好！", "你好！有什么可以帮助你的吗？", session1_id)
    memory.add_turn("今天天气如何？", "今天天气晴朗。", session1_id)

    # 开始新会话 2
    session2_id = memory.start_new_session()
    memory.add_turn("讲个笑话吧。", "为什么电脑去医院？因为它需要‘重启’一下！", session2_id)

    # 切换回会话 1 并继续
    memory.set_current_session(session1_id)
    memory.add_turn("谢谢！", "不客气！", session1_id)

    print("--- 所有会话列表 ---")
    print(memory.get_sessions_list())

    print(f"\n--- 会话 {session1_id} 的历史 ---")
    print(memory.get_history(session1_id))

    print(f"\n--- 会话 {session1_id} 的上下文 (最近2轮) ---")
    print(memory.get_context(session1_id, last_n=2))

    print(f"\n--- 会话 {session2_id} 的历史 ---")
    print(memory.get_history(session2_id))




